# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import os
import sys
import time
import json
import asyncio
import logging
import importlib
from typing import List, Dict
from datetime import datetime, timedelta
from fastapi import FastAPI
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.interval import IntervalTrigger
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.date import DateTrigger
from config import get_settings
from common.cache.queue_cache import QueueCache
from common.enums.apply import MusicEnum
from common.enums.apply import VideoEnum
from common.enums.apply import DrawsEnum
from common.enums.agent import AgentEnum
from common.enums.public import CrontabEnum
from common.utils.config import ConfigUtil
from common.utils.cache import RedisUtil
from common.utils.times import TimeUtil
from common.models.sys import SysCrontabModel
from common.models.music import MusicRecordsModel
from common.models.video import VideoRecordsModel
from common.models.draws import DrawsRecordsModel
from common.models.agents import AgentKnowArchiveModel


scheduler = AsyncIOScheduler()
logger = logging.getLogger(__name__)


class AppEvents:
    @classmethod
    async def startup(cls, _app: FastAPI):
        with open("scheduler.pid", "w") as f:
            f.write(str(os.getpid()))

        delay = 1 if get_settings().SERVER_WORKERS == 1 else 15
        asyncio.create_task(cls._redis_subscribe())
        run_date = datetime.now() + timedelta(seconds=delay)
        scheduler.add_job(cls._init_crontab, DateTrigger(run_date=run_date))
        scheduler.start()

    @classmethod
    async def shutdown(cls, _app: FastAPI):
        scheduler.shutdown()

    @classmethod
    async def _init_crontab(cls):
        """ Initialize scheduled """
        with open("scheduler.pid", "r") as f:
            pid = f.read()
            if pid != str(os.getpid()):
                return False

        await ConfigUtil.set("sys", "process_id", pid)

        tasks = []
        enums = {1: "Success", 2: "Stop", 3: "Error"}
        crontabs = await SysCrontabModel.filter(is_delete=0).order_by("-id").all()
        for crontab in crontabs:
            task = [crontab.id, crontab.name, crontab.command, crontab.concurrent]
            if crontab.status != CrontabEnum.CRON_ING:
                continue

            try:
                module = importlib.import_module(crontab.command)
                clz = getattr(module, "Command", None)
                if not clz:
                    raise AttributeError(f"Cron: '{crontab.command}' has no 'Command' class")
                fun = getattr(clz, "execute", None)
                if not fun:
                    raise AttributeError(f"Cron: '{crontab.command}' has no 'execute' method")
            except ModuleNotFoundError:
                wrong: str = f"Crontab 'module' not found '{crontab.command}'"
                logger.error(wrong)
                crontab.error = wrong
                crontab.status = CrontabEnum.CRON_ERROR
                tasks.append(task + [crontab.status])
                await crontab.save()
                continue
            except AttributeError as e:
                if crontab.command in sys.modules:
                    del sys.modules[crontab.command]
                logger.error(str(e))
                crontab.error = str(e)
                crontab.status = CrontabEnum.CRON_ERROR
                tasks.append(task + [crontab.status])
                await crontab.save()
                continue

            rules: List[dict] = json.loads(crontab.rules or "[]")
            params: Dict[str, any] = json.loads(crontab.params or "{}")

            condition = {}
            for item in rules:
                if crontab.trigger == "interval" and item.get("key") not in ["start_date", "end_date"]:
                    condition[item.get("key")] = int(item.get("value"))
                else:
                    condition[item.get("key")] = item.get("value")

            _trigger_fun = None
            if crontab.trigger == "interval":
                _trigger_fun = IntervalTrigger(**condition)
            elif crontab.trigger == "cron":
                _trigger_fun = CronTrigger(**condition)
            elif crontab.trigger == "date":
                _trigger_fun = DateTrigger(**condition)

            tasks.append(task + [crontab.status])
            for i in range(crontab.concurrent):
                job = crontab.command + "." + str(i + 1)
                params["w_id"] = crontab.id
                params["w_ix"] = int(i + 1)
                params["w_pid"] = int(os.getpid())
                params["w_job"] = job
                scheduler.add_job(fun, _trigger_fun, id=job, name=crontab.name, kwargs=params)

        if os.path.exists("./banner.txt"):
            startup_time = TimeUtil.timestamp_to_date(int(time.time()))
            with open("./banner.txt", "r") as f:
                template = f.read()
                template = template.replace("{{startup_time}}", startup_time)
                template = template.replace("{{process_id}}", str(os.getpid()))
                print(template)

        for job in tasks:
            job[4] = enums[job[4]]
            print(job)

        print("-" * 57)
        await cls._reset_queue()
        return True

    @classmethod
    async def _redis_subscribe(cls):
        """ Redis/Push/Sub"""
        try:
            pubsub = await RedisUtil.subscribe("topical")
            async for message in pubsub.listen():
                if message["type"] != "message":
                    continue

                event = json.loads(str(message["data"] or "{}"))
                scene = event.get("scene")
                if scene == "cron":
                    from apps.admin.service.system.crontab_service import CrontabService
                    await CrontabService.cron_subscribe(str(event["data"]))
        except Exception as e:
            print(str(e))

    @classmethod
    async def _reset_queue(cls):
        # 向量队列重置
        await QueueCache.queue_del(QueueCache.EM_JOB)
        embed_lists = await (AgentKnowArchiveModel
                             .filter(train_mode="chunk")
                             .filter(status__in=[AgentEnum.EMB_LINE, AgentEnum.EMB_ING])
                             .filter(is_delete=0)
                             .values("id"))
        for item in embed_lists:
            await QueueCache.queue_push(QueueCache.EM_JOB, str(item["id"]))

        # 拆分队列重置
        await QueueCache.queue_del(QueueCache.QA_JOB)
        split_lists = await (AgentKnowArchiveModel
                             .filter(train_mode="qa")
                             .filter(qa_status__in=[AgentEnum.EMB_LINE, AgentEnum.EMB_ING])
                             .filter(is_delete=0)
                             .values("id"))
        for item in split_lists:
            await QueueCache.queue_push(QueueCache.QA_JOB, str(item["id"]))

        # 绘画队列重置
        await QueueCache.queue_del(QueueCache.DRAW_JOB)
        draws_lists = await (DrawsRecordsModel.filter(is_delete=0, status=DrawsEnum.STATUS_ING).values("id"))
        for item in draws_lists:
            await QueueCache.queue_push(QueueCache.DRAW_JOB, item["id"])

        # 视频队列重置
        await QueueCache.queue_del(QueueCache.VIDEO_JOB)
        video_lists = await (VideoRecordsModel.filter(is_delete=0, status=VideoEnum.STATUS_ING).values("id"))
        for item in video_lists:
            await QueueCache.queue_push(QueueCache.VIDEO_JOB, item["id"])

        # 音乐队列重置
        await QueueCache.queue_del(QueueCache.VIDEO_JOB)
        music_lists = await (MusicRecordsModel.filter(is_delete=0, status=MusicEnum.STATUS_ING).values("id"))
        for item in music_lists:
            await QueueCache.queue_push(QueueCache.MUSIC_JOB, item["id"])
